import tensorflow as tf

# 加载MNIST数据集
mnist = tf.keras.datasets.mnist
(x_train, y_train), (x_test, y_test) = mnist.load_data()

# 改变维度来适应CNN网络的输入要求
x_train4D = x_train.reshape(x_train.shape[0],28,28,1).astype('float32') 
x_test4D = x_test.reshape(x_test.shape[0],28,28,1).astype('float32')

# 数据预处理 
# *********** 请介绍数据预处理步骤 **************
x_train, x_test = x_train4D / 255.0, x_test4D / 255.0

# 构建卷积神经网络模型
model = tf.keras.models.Sequential([ 
    # 第一层卷积层：32个3x3卷积核，使用ReLU激活函数，padding='same'保持输出与输入大小一致
    tf.keras.layers.Conv2D(filters=32, kernel_size=(3, 3), padding='same',
                           input_shape=(28, 28, 1), activation='relu'),
    
    # 第一层池化层：使用2x2的池化窗口进行最大池化，缩小特征图的尺寸
    tf.keras.layers.MaxPooling2D(pool_size=(2, 2)),
    
    # 第二层卷积层：64个3x3卷积核，同样使用ReLU激活函数，padding='same'保持输出大小一致
    tf.keras.layers.Conv2D(filters=64, kernel_size=(3, 3), padding='same', activation='relu'),
    
    # 第二层池化层：同样使用2x2池化窗口进行最大池化
    tf.keras.layers.MaxPooling2D(pool_size=(2, 2)),
    
    # Dropout层：随机丢弃25%的神经元，防止过拟合
    tf.keras.layers.Dropout(0.25),
    
    # Flatten层：将多维的卷积层输出展平成一维，用于输入到全连接层
    tf.keras.layers.Flatten(),
    
    # 第一个全连接层：128个神经元，ReLU激活函数，学习到更复杂的特征
    tf.keras.layers.Dense(128, activation='relu'),
    
    # Dropout层：随机丢弃50%的神经元，进一步防止过拟合
    tf.keras.layers.Dropout(0.5),
    
    # 输出层：10个神经元，代表数字0-9的类别，使用Softmax激活函数，输出各类的概率
    tf.keras.layers.Dense(10, activation='softmax')
])
# 定义交叉熵损失函数
loss_fn = tf.keras.losses.SparseCategoricalCrossentropy(from_logits=True)

# 定义优化器和学习率
optimizer = tf.keras.optimizers.Adam(learning_rate=0.001)

# 编译模型
model.compile(optimizer=optimizer,
              loss=loss_fn,
              metrics=['accuracy'])

# 训练模型
model.fit(x_train, y_train, epochs=5, validation_data=(x_test, y_test))

# 评估模型
model.evaluate(x_test, y_test)
